#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RedFire PostgreSQL 分区自动化服务
================================

服务功能:
- 定时创建和维护分区
- 分区性能监控
- 自动优化建议
- Web API管理接口

Author: RedFire Team
Date: 2025-09-16
Version: v1.0.0
"""

import asyncio
import logging
import os
import sys
import signal
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional
import yaml
import json
from dataclasses import asdict
import schedule
import uvicorn
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
import psutil

# 添加项目路径
current_dir = Path(__file__).parent
project_root = current_dir.parent.parent
sys.path.insert(0, str(project_root))

from databases.partition.postgresql_partition_manager import (
    PostgreSQLPartitionManager,
    PartitionConfig,
    VNPY_PARTITION_CONFIGS
)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class PartitionService:
    """分区管理服务"""
    
    def __init__(self, config_path: str = None):
        self.config_path = config_path or os.path.join(current_dir, "partition_config.yaml")
        self.config = None
        self.partition_manager = None
        self.running = False
        self.scheduler_task = None
        self.stats = {
            'service_start_time': None,
            'last_maintenance_time': None,
            'partitions_created': 0,
            'partitions_dropped': 0,
            'maintenance_cycles': 0,
            'errors': 0
        }
    
    async def load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config_data = yaml.safe_load(f)
            
            # 环境变量替换
            config_data = self._replace_env_vars(config_data)
            
            # 获取环境特定配置
            env = os.getenv('REDFIRE_ENV', 'development')
            if env in config_data.get('environments', {}):
                env_config = config_data['environments'][env]
                config_data = self._merge_config(config_data, env_config)
            
            self.config = config_data
            logger.info(f"Partition service configuration loaded (env: {env})")
            return config_data
            
        except Exception as e:
            logger.error(f"Failed to load config: {e}")
            raise
    
    def _replace_env_vars(self, data):
        """递归替换环境变量"""
        if isinstance(data, dict):
            return {k: self._replace_env_vars(v) for k, v in data.items()}
        elif isinstance(data, list):
            return [self._replace_env_vars(item) for item in data]
        elif isinstance(data, str) and data.startswith('${') and data.endswith('}'):
            env_expr = data[2:-1]
            if ':-' in env_expr:
                var_name, default_value = env_expr.split(':-', 1)
                return os.getenv(var_name, default_value)
            else:
                return os.getenv(env_expr, data)
        else:
            return data
    
    def _merge_config(self, base_config: Dict, env_config: Dict) -> Dict:
        """合并环境特定配置"""
        result = base_config.copy()
        
        def deep_merge(target, source):
            for key, value in source.items():
                if key in target and isinstance(target[key], dict) and isinstance(value, dict):
                    deep_merge(target[key], value)
                else:
                    target[key] = value
        
        deep_merge(result, env_config)
        return result
    
    async def initialize(self):
        """初始化分区服务"""
        await self.load_config()
        
        # 创建分区管理器
        self.partition_manager = PostgreSQLPartitionManager(self.config['database'])
        await self.partition_manager.initialize()
        
        # 注册表分区配置
        for table_name, table_config in self.config['table_partitions'].items():
            if table_config.get('enabled', False):
                partition_config = PartitionConfig(
                    table_name=table_name,
                    partition_column=table_config['partition_column'],
                    partition_type=table_config['partition_type'],
                    retention_months=table_config['retention_months'],
                    auto_create_future=table_config['auto_create_future'],
                    compression_enabled=table_config.get('compression_enabled', False),
                    compression_after_days=table_config.get('compression_after_days', 7),
                    create_indexes=table_config.get('create_indexes', True),
                    index_patterns=table_config.get('index_patterns', [])
                )
                
                await self.partition_manager.register_table(partition_config)
        
        logger.info("Partition service initialized successfully")
    
    async def start_service(self):
        """启动分区服务"""
        if self.running:
            logger.warning("Partition service is already running")
            return
        
        self.running = True
        self.stats['service_start_time'] = datetime.now()
        
        logger.info("Starting partition management service...")
        
        # 执行初始维护
        await self.perform_initial_maintenance()
        
        # 启动定时任务调度器
        self.scheduler_task = asyncio.create_task(self._scheduler_loop())
        
        logger.info("Partition service started successfully")
    
    async def perform_initial_maintenance(self):
        """执行初始维护"""
        logger.info("Performing initial partition maintenance...")
        
        try:
            # 执行自动维护
            maintenance_results = await self.partition_manager.auto_maintain_all_tables()
            
            # 更新统计信息
            for table_name, result in maintenance_results.items():
                if result.get('status') == 'success':
                    self.stats['partitions_created'] += result.get('future_partitions_created', 0)
                    self.stats['partitions_dropped'] += result.get('old_partitions_dropped', 0)
                else:
                    self.stats['errors'] += 1
            
            self.stats['maintenance_cycles'] += 1
            self.stats['last_maintenance_time'] = datetime.now()
            
            logger.info("Initial maintenance completed")
            logger.info(f"Maintenance results: {maintenance_results}")
            
        except Exception as e:
            logger.error(f"Initial maintenance failed: {e}")
            self.stats['errors'] += 1
    
    async def _scheduler_loop(self):
        """定时任务调度循环"""
        maintenance_tasks = self.config.get('maintenance_tasks', {})
        interval_hours = self.config['global_partition_settings']['maintenance_interval_hours']
        
        while self.running:
            try:
                # 检查是否需要执行维护
                if self._should_run_maintenance():
                    await self._run_scheduled_maintenance()
                
                # 等待下一个检查周期
                await asyncio.sleep(3600)  # 每小时检查一次
                
            except asyncio.CancelledError:
                logger.info("Scheduler loop cancelled")
                break
            except Exception as e:
                logger.error(f"Error in scheduler loop: {e}")
                await asyncio.sleep(3600)
    
    def _should_run_maintenance(self) -> bool:
        """检查是否应该运行维护"""
        if not self.stats['last_maintenance_time']:
            return True
        
        interval_hours = self.config['global_partition_settings']['maintenance_interval_hours']
        next_maintenance = self.stats['last_maintenance_time'] + timedelta(hours=interval_hours)
        
        return datetime.now() >= next_maintenance
    
    async def _run_scheduled_maintenance(self):
        """运行定时维护"""
        logger.info("Running scheduled partition maintenance...")
        
        try:
            # 执行自动维护
            maintenance_results = await self.partition_manager.auto_maintain_all_tables()
            
            # 更新统计信息
            for table_name, result in maintenance_results.items():
                if result.get('status') == 'success':
                    self.stats['partitions_created'] += result.get('future_partitions_created', 0)
                    self.stats['partitions_dropped'] += result.get('old_partitions_dropped', 0)
                else:
                    self.stats['errors'] += 1
            
            self.stats['maintenance_cycles'] += 1
            self.stats['last_maintenance_time'] = datetime.now()
            
            logger.info("Scheduled maintenance completed")
            
        except Exception as e:
            logger.error(f"Scheduled maintenance failed: {e}")
            self.stats['errors'] += 1
    
    async def stop_service(self):
        """停止分区服务"""
        if not self.running:
            return
        
        logger.info("Stopping partition service...")
        self.running = False
        
        if self.scheduler_task:
            self.scheduler_task.cancel()
            try:
                await self.scheduler_task
            except asyncio.CancelledError:
                pass
        
        if self.partition_manager:
            await self.partition_manager.cleanup()
        
        logger.info("Partition service stopped")
    
    async def get_service_status(self) -> Dict[str, Any]:
        """获取服务状态"""
        status = {
            'running': self.running,
            'stats': self.stats.copy(),
            'config_loaded': self.config is not None,
            'registered_tables': list(self.partition_manager.partition_configs.keys()) if self.partition_manager else [],
            'system_info': {
                'cpu_percent': psutil.cpu_percent(),
                'memory_percent': psutil.virtual_memory().percent,
                'disk_percent': psutil.disk_usage('/').percent,
                'process_id': os.getpid()
            }
        }
        
        if self.stats['service_start_time']:
            uptime = datetime.now() - self.stats['service_start_time']
            status['uptime_seconds'] = uptime.total_seconds()
        
        return status
    
    async def get_partition_summary(self) -> Dict[str, Any]:
        """获取分区摘要信息"""
        if not self.partition_manager:
            raise RuntimeError("Partition manager not initialized")
        
        summary = {}
        
        for table_name in self.partition_manager.partition_configs.keys():
            try:
                partitions_info = await self.partition_manager.get_partition_info(table_name)
                
                summary[table_name] = {
                    'total_partitions': len(partitions_info),
                    'total_size_bytes': sum(p.size_bytes for p in partitions_info),
                    'total_rows': sum(p.row_count for p in partitions_info),
                    'date_range': {
                        'earliest': min(p.start_date for p in partitions_info).isoformat() if partitions_info else None,
                        'latest': max(p.end_date for p in partitions_info).isoformat() if partitions_info else None
                    },
                    'partitions': [
                        {
                            'name': p.partition_name,
                            'start_date': p.start_date.isoformat(),
                            'end_date': p.end_date.isoformat(),
                            'size_mb': p.size_bytes // (1024 * 1024),
                            'row_count': p.row_count
                        }
                        for p in partitions_info
                    ]
                }
                
            except Exception as e:
                logger.error(f"Error getting partition info for {table_name}: {e}")
                summary[table_name] = {'error': str(e)}
        
        return summary
    
    async def force_maintenance(self, table_name: str = None) -> Dict[str, Any]:
        """强制执行维护"""
        if not self.partition_manager:
            raise RuntimeError("Partition manager not initialized")
        
        if table_name:
            # 维护指定表
            if table_name not in self.partition_manager.partition_configs:
                raise ValueError(f"Table {table_name} not registered")
            
            # 创建未来分区
            config = self.partition_manager.partition_configs[table_name]
            future_partitions = await self.partition_manager.create_future_partitions(
                table_name, config.auto_create_future
            )
            
            # 清理过期分区
            dropped_partitions = await self.partition_manager.cleanup_old_partitions(table_name)
            
            # 优化分区
            optimization_results = await self.partition_manager.optimize_partitions(table_name)
            
            return {
                table_name: {
                    'future_partitions_created': len(future_partitions),
                    'old_partitions_dropped': len(dropped_partitions),
                    'optimization_results': optimization_results
                }
            }
        else:
            # 维护所有表
            return await self.partition_manager.auto_maintain_all_tables()
    
    async def create_partition_manually(self, table_name: str, start_date: str, end_date: str) -> str:
        """手动创建分区"""
        if not self.partition_manager:
            raise RuntimeError("Partition manager not initialized")
        
        from datetime import datetime
        start_dt = datetime.fromisoformat(start_date).date()
        end_dt = datetime.fromisoformat(end_date).date()
        
        partition_name = await self.partition_manager.create_partition(table_name, start_dt, end_dt)
        
        # 更新统计
        self.stats['partitions_created'] += 1
        
        return partition_name


# 创建全局服务实例
partition_service = PartitionService()

# 创建FastAPI应用
app = FastAPI(
    title="RedFire PostgreSQL Partition Management Service",
    description="PostgreSQL分区自动化管理服务",
    version="1.0.0"
)


@app.on_event("startup")
async def startup_event():
    """服务启动事件"""
    try:
        await partition_service.initialize()
        logger.info("Partition service API started")
    except Exception as e:
        logger.error(f"Failed to start partition service: {e}")
        raise


@app.on_event("shutdown")
async def shutdown_event():
    """服务关闭事件"""
    await partition_service.stop_service()
    logger.info("Partition service API shutdown")


@app.get("/health")
async def health_check():
    """健康检查"""
    try:
        status = await partition_service.get_service_status()
        return JSONResponse(
            status_code=200,
            content={
                "status": "healthy",
                "timestamp": datetime.now().isoformat(),
                "service": status
            }
        )
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={
                "status": "unhealthy",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
        )


@app.get("/status")
async def get_status():
    """获取详细状态"""
    status = await partition_service.get_service_status()
    return JSONResponse(content=status)


@app.get("/partitions")
async def get_partition_summary():
    """获取分区摘要"""
    try:
        summary = await partition_service.get_partition_summary()
        return JSONResponse(content=summary)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/partitions/{table_name}")
async def get_table_partitions(table_name: str):
    """获取指定表的分区信息"""
    try:
        if not partition_service.partition_manager:
            raise HTTPException(status_code=503, detail="Partition manager not initialized")
        
        partitions_info = await partition_service.partition_manager.get_partition_info(table_name)
        return JSONResponse(content=[asdict(p) for p in partitions_info])
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/service/start")
async def start_service(background_tasks: BackgroundTasks):
    """启动分区服务"""
    try:
        background_tasks.add_task(partition_service.start_service)
        return JSONResponse(content={"message": "Partition service starting"})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/service/stop")
async def stop_service():
    """停止分区服务"""
    try:
        await partition_service.stop_service()
        return JSONResponse(content={"message": "Partition service stopped"})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/maintenance/run")
async def run_maintenance(table_name: Optional[str] = None):
    """运行维护任务"""
    try:
        result = await partition_service.force_maintenance(table_name)
        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/partitions/{table_name}/create")
async def create_partition(table_name: str, start_date: str, end_date: str):
    """手动创建分区"""
    try:
        partition_name = await partition_service.create_partition_manually(
            table_name, start_date, end_date
        )
        return JSONResponse(content={
            "message": "Partition created successfully",
            "partition_name": partition_name
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/recommendations/{table_name}")
async def get_maintenance_recommendations(table_name: str):
    """获取维护建议"""
    try:
        if not partition_service.partition_manager:
            raise HTTPException(status_code=503, detail="Partition manager not initialized")
        
        recommendations = await partition_service.partition_manager.get_maintenance_recommendations(table_name)
        return JSONResponse(content=recommendations)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/config")
async def get_config():
    """获取当前配置"""
    if not partition_service.config:
        raise HTTPException(status_code=404, detail="Configuration not loaded")
    
    # 过滤敏感信息
    safe_config = partition_service.config.copy()
    if 'database' in safe_config:
        safe_config['database']['password'] = "***"
    
    return JSONResponse(content=safe_config)


def setup_signal_handlers():
    """设置信号处理器"""
    def signal_handler(signum, frame):
        logger.info(f"Received signal {signum}")
        asyncio.create_task(partition_service.stop_service())
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)


async def run_service():
    """运行分区服务"""
    setup_signal_handlers()
    
    try:
        await partition_service.initialize()
        await partition_service.start_service()
        
        # 保持服务运行
        while partition_service.running:
            await asyncio.sleep(1)
            
    except KeyboardInterrupt:
        logger.info("Received keyboard interrupt")
    except Exception as e:
        logger.error(f"Service error: {e}")
    finally:
        await partition_service.stop_service()


def run_api_server():
    """运行API服务器"""
    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=8002,
        log_level="info"
    )
    server = uvicorn.Server(config)
    asyncio.run(server.serve())


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="RedFire PostgreSQL Partition Management Service")
    parser.add_argument(
        "--mode", 
        choices=["service", "api", "once", "init"], 
        default="service",
        help="运行模式: service(服务模式), api(API模式), once(单次维护), init(初始化)"
    )
    parser.add_argument(
        "--config", 
        help="配置文件路径",
        default=None
    )
    parser.add_argument(
        "--table",
        help="指定表名(仅在once模式下有效)",
        default=None
    )
    
    args = parser.parse_args()
    
    if args.config:
        partition_service.config_path = args.config
    
    if args.mode == "service":
        # 服务模式：持续运行
        asyncio.run(run_service())
    elif args.mode == "api":
        # API模式：启动Web API服务
        run_api_server()
    elif args.mode == "once":
        # 单次模式：执行一次维护后退出
        async def run_once():
            await partition_service.initialize()
            result = await partition_service.force_maintenance(args.table)
            print(json.dumps(result, indent=2, default=str))
            await partition_service.stop_service()
        
        asyncio.run(run_once())
    elif args.mode == "init":
        # 初始化模式：设置分区表后退出
        async def run_init():
            await partition_service.initialize()
            await partition_service.perform_initial_maintenance()
            await partition_service.stop_service()
            print("Partition initialization completed")
        
        asyncio.run(run_init())
